from langchain.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from typing import List
import os
import yaml


class DocumentProcessor:
    """文档处理器"""

    def __init__(self, config_path: str = "config.yaml"):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)

        rag_config = config['rag']
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=rag_config['chunk_size'],
            chunk_overlap=rag_config['chunk_overlap'],
            length_function=len,
        )

    def load_documents(self, directory: str) -> List[Document]:
        """加载目录中的所有文档"""
        documents = []

        for filename in os.listdir(directory):
            filepath = os.path.join(directory, filename)

            if filename.endswith('.pdf'):
                loader = PyPDFLoader(filepath)
                docs = loader.load()
            elif filename.endswith('.txt'):
                loader = TextLoader(filepath, encoding='utf-8')
                docs = loader.load()
            else:
                continue

            # 分割文档
            split_docs = self.text_splitter.split_documents(docs)
            documents.extend(split_docs)

        return documents